//
// Created by 86155 on 2025/4/16.
//

#include "FeISelLowering.h"
#include "MCTargetDesc/FeMCTargetDesc.h"
#include "llvm/CodeGen/CallingConvLower.h"
#include "FeSubtarget.h"
#include "FeTargetMachine.h"
#include <unordered_map> // 用于哈希表
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "FeDataCtrl.h"
#include <string>
using namespace llvm;
#include "FeGenCallingConv.inc"

// 1. 定义全局映射表（线程安全需自行处理，此处假设单线程编译场景）
struct StoreMetadata {
  SDValue StoreNode;   // StoreSDNode 节点
  SDValue Chain;       // Store 节点的依赖链（Chain）
  SDValue Value;       // Store 节点的值 存入帧
  SDLoc DL;            // Store 节点的位置信息（可选）
};

std::unordered_map<int, StoreMetadata> FrameIndexToLastStore;

// 2. 辅助函数：判断 SDValue 是否为 FrameIndexSDNode 并提取索引
bool isFrameIndexSDValue(SDValue V, int &FI) {
  if (FrameIndexSDNode *FIN = dyn_cast<FrameIndexSDNode>(V)) {
    FI = FIN->getIndex();
    return true;
  }
  return false;
}

// 3. 记录 Store 操作到全局映射表
void recordStoreOperation(SDValue StoreNode, SDValue Chain, SDValue Value, SDLoc DL) {
  int FI;
  if (!isFrameIndexSDValue(StoreNode->getOperand(2), FI)) { // Store 的Operand 2 是地址
    return; // 非 FrameIndex 地址，不处理
  }

  StoreMetadata Meta;
  Meta.StoreNode = StoreNode;
  Meta.Chain = Chain;
  Meta.Value = Value;
  Meta.DL = DL;

  FrameIndexToLastStore[FI] = Meta; // 覆盖旧记录，保留最后一次 Store
  printf( "Recorded Store for FrameIndex %d\n",FI);
}

// 4. 查询最后一次 Store 操作
std::optional<StoreMetadata> getLastStoreForFrameIndex(int FI) {
  auto It = FrameIndexToLastStore.find(FI);
  if (It != FrameIndexToLastStore.end()) {
    return It->second;
  }
  return std::nullopt;
}

FeTargetLowering::FeTargetLowering(const TargetMachine &TM,
                                     const FeSubtarget &STI)
    : TargetLowering(TM), Subtarget(STI) {
  /// 注册RegiserClass
  /// 还要处理合法化（类型和操作）
  addRegisterClass(MVT::i32, &Fe::GPRRegClass);

  setOperationAction(ISD::LOAD, MVT::i32, Custom);
  setOperationAction(ISD::STORE, MVT::i32, Custom);
  //对BR_CC阶段进行展开
  setOperationAction(ISD::BR_CC, MVT::i32, Expand);
  computeRegisterProperties(STI.getRegisterInfo());
}

SDValue FeTargetLowering::LowerFormalArguments(
    SDValue Chain, CallingConv::ID CallConv, bool IsVarArg,
    const SmallVectorImpl<ISD::InputArg> &Ins, const SDLoc &DL,
    SelectionDAG &DAG, SmallVectorImpl<SDValue> &InVals) const {

  MachineFunction &MF = DAG.getMachineFunction(); //函数
  MachineFrameInfo &MFI = MF.getFrameInfo();

  SmallVector<CCValAssign, 16> ArgLocs;
  CCState CCInfo(CallConv, IsVarArg, MF, ArgLocs, *DAG.getContext());
  CCInfo.AnalyzeFormalArguments(Ins, CC_Fe);

  SDValue ArgValue;
  for (unsigned i = 0, e = ArgLocs.size(); i != e; ++i) {
    CCValAssign &VA = ArgLocs[i];
    if (VA.isRegLoc()) {
      MVT RegVT = VA.getLocVT();
      Register Reg = MF.addLiveIn(VA.getLocReg(), &Fe::GPRRegClass);
      ArgValue = DAG.getCopyFromReg(Chain, DL, Reg, RegVT);
      InVals.push_back(ArgValue);
    } else {//处理形参 li
      assert(VA.isMemLoc());
      MVT ValVT = VA.getValVT();
      int Offset = VA.getLocMemOffset();
      int FI = MFI.CreateFixedObject(ValVT.getSizeInBits() / 8, Offset, true);
      SDValue FIN = DAG.getFrameIndex(FI, getPointerTy(DAG.getDataLayout()));
      SDValue Val = DAG.getLoad(
          ValVT, DL, Chain, FIN,
          MachinePointerInfo::getFixedStack(DAG.getMachineFunction(), FI));
      InVals.push_back(Val);
    }
  }
  return Chain;
}
//固定模版 OutVals拷贝到物理寄存器  通过调用约定决策CallConv
SDValue FeTargetLowering::LowerReturn(SDValue Chain, CallingConv::ID CallConv,
                               bool IsVarArg,
                               const SmallVectorImpl<ISD::OutputArg> &Outs/*输出参数*/,
                               const SmallVectorImpl<SDValue> &OutVals/*输出值*/,
                               const SDLoc &DL, SelectionDAG &DAG) const {
  /// 1. 返回物理寄存器（调用约定的限值）
  SmallVector<CCValAssign, 16> RVLocs;   //分析哪些寄存器可以用来做返回值
  CCState CCInfo(CallConv, IsVarArg, DAG.getMachineFunction(), RVLocs,
                 *DAG.getContext()/*上下文*/);
  CCInfo.AnalyzeReturn(Outs, RetCC_Fe);   //分析返回值  RVLocs目前已有信息
  SDValue Glue;
  SmallVector<SDValue, 4> RetOps(1, Chain);
  for (unsigned i = 0, e = RVLocs.size(); i < e; ++i) {
    CCValAssign &VA = RVLocs[i];
    assert(VA.isRegLoc() && "Can only return in registers!");
    //创建节点
    Chain = DAG.getCopyToReg(Chain, DL/*位置信息*/, VA.getLocReg(), OutVals[i], Glue);
    Glue = Chain.getValue(1);   //设置成第一个输出值
    RetOps.push_back(DAG.getRegister(VA.getLocReg(), VA.getLocVT()));
  }
  RetOps[0] = Chain;
  if (Glue.getNode()) {
    RetOps.push_back(Glue);
  }
  return DAG.getNode(FeISD::RET_GLUE, DL, MVT::Other, RetOps);
}

SDValue FeTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
  switch (Op.getOpcode()) {
  case ISD::LOAD:
  {
//    return LowerLOAD(Op, DAG);
    return SDValue();
  }
  case ISD::STORE:
  {
//    return LowerSTORE(Op, DAG);
    return SDValue();
  }
  default:
    llvm::llvm_unreachable_internal("unknown op");
  }
  return SDValue();
}

/*
 * 1. 消除store节点
 * 2. 将load替换为addi
 */
SDValue FeTargetLowering::LowerLOAD(SDValue Op,SelectionDAG &DAG) const {
  EVT VT = Op.getValueType();
  LoadSDNode *N = dyn_cast<LoadSDNode>(Op);
  SDLoc DL(N);

  // 获取三个核心操作数
  SDValue Chain = N->getOperand(0);      // 依赖链 (例如t5的ch输出)
  SDValue Address = N->getOperand(1);    // 源地址 (例如t2: FrameIndex<0>)
      
  // 打印操作数信息
  printf("Load操作数分析:\n");
  printf("  Operand 0 (Chain): \n");
//  Chain.getNode()->dump();
      
  printf("  Operand 1 (Address): \n");
  if (FrameIndexSDNode *FIN = dyn_cast<FrameIndexSDNode>(Address)) {
    printf("FrameIndex = %d\n" ,FIN->getIndex());
  } else {
//    Address.getNode()->dump();
  }

  int FI;
  if (!isFrameIndexSDValue(Address, FI)) {
    // 非 FrameIndex 地址，按原始逻辑处理
    printf("非 FrameIndex 地址，按原始逻辑处理");
  }

  // 查询全局映射表，获取最后一次 Store
  if (auto OptMeta = getLastStoreForFrameIndex(FI)) {
//    OptMeta.value().Value.getNode()->dump();

    // 建立跨 DAG 的依赖：将当前 Load 的 Chain 与 Store 的 Chain 合并
    Chain = DAG.getNode(ISD::TokenFactor, DL, MVT::Other, Chain, OptMeta.value().Chain);
    printf( "Established cross-DAG dependency for FrameIndex %d\n",FI);
  }


  SDValue Alignment;
  // 检查对齐信息 (如果有)
  if (N->getNumOperands() > 2) {
    Alignment = N->getOperand(2);
    printf("  Operand 2 (Alignment): \n");
//    Alignment.getNode()->dump();
  }

  printf("\n");
//  return DAG.getNode(FeISD::LOA,DL,VT,Chain,Address,Alignment);
  return SDValue();
}

SDValue FeTargetLowering::LowerSTORE(SDValue Op,SelectionDAG &DAG) const{
  EVT VT = Op.getValueType();
  StoreSDNode *N = dyn_cast<StoreSDNode>(Op);
  SDLoc DL(N);
  // 获取三个核心操作数
  SDValue Chain = N->getOperand(0);     // 依赖链 (例如t0: EntryToken)
  SDValue Value = N->getOperand(1);     // 存储的值 (例如t1: Constant<10>)
  SDValue Address = N->getOperand(2);   // 目标地址 (例如t2: FrameIndex<0>)

  // 打印操作数信息
  printf("Store操作数分析:\n");
  printf("  Operand 0 (Chain): \n");
//  Chain.getNode()->dump();

  // 验证与图中节点的连接关系
  if (Chain.getNode() && Chain.getNode()->getOpcode() == ISD::EntryToken) {
    printf( "✓ 依赖链连接到EntryToken (如t0)\n");
  }

  APInt ConstValue;
  printf("  Operand 1 (Value): \n");
  if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(Value)) {
    ConstValue = CN->getAPIntValue();
    printf("Constant Value = %d\n",ConstValue.getSExtValue());       // Constant Value = 10
  } else {
//    Value.getNode()->dump();                                                // t7: i32 = add nsw t6, Constant:i32<10>
  }

  printf("  Operand 2 (Address):\n ");
  if (FrameIndexSDNode *FIN = dyn_cast<FrameIndexSDNode>(Address)) {
    printf("FrameIndex = %d\n",FIN->getIndex());                     // FrameIndex = 0
  } else {
//    Address.getNode()->dump();i
  }

  SDValue Alignment;
  // 检查对齐信息
  if (N->getNumOperands() > 3) {
    Alignment = N->getOperand(3);
    printf("  Operand 3 (Alignment): \n");
//    Alignment.getNode()->dump();
  }
  recordStoreOperation(Op, Chain, Value, DL);
  printf("\n");
//  return DAG.getNode(FeISD::STO,DL,VT,Chain,Value,Address,Alignment);
  return SDValue();
}

//调用处
SDValue FeTargetLowering::LowerCall(CallLoweringInfo &CLI,
                                     SmallVectorImpl<SDValue> &InVals) const {
  SelectionDAG &DAG = CLI.DAG;//DAG数据结构
  SDLoc &DL = CLI.DL;         //所在位置
  SmallVectorImpl<ISD::OutputArg> &Outs = CLI.Outs;//形参
  SmallVectorImpl<SDValue> &OutVals = CLI.OutVals; //实参
  SmallVectorImpl<ISD::InputArg> &Ins = CLI.Ins;   //返回值
  SDValue Chain = CLI.Chain;                       //控制链
  SDValue Callee = CLI.Callee;                     //节点
  CallingConv::ID CallConv = CLI.CallConv;         //调用约定
  bool IsVarArg = CLI.IsVarArg;                    //参数可变

  MachineFunction &MF = DAG.getMachineFunction();
  MachineFrameInfo &MFI = MF.getFrameInfo();
  SmallVector<CCValAssign, 16> ArgLocs;
  CCState CCInfo(CallConv, IsVarArg, MF, ArgLocs, *DAG.getContext());
  CCInfo.AnalyzeCallOperands(Outs, CC_Fe);

  //记录物理寄存器
  SmallVector<std::pair<unsigned, SDValue>> RegsPairs;
  SDValue StackPtr;

  //实参处理
  for (unsigned i = 0, e = ArgLocs.size(); i != e; ++i) {
    CCValAssign &VA = ArgLocs[i];
    if(VA.isRegLoc()){
      RegsPairs.push_back(std::make_pair(VA.getLocReg(), OutVals[i]));
    }
    else{//寄存器不够时
      // store
      assert(VA.isMemLoc());
      if (!StackPtr.getNode()) {
        StackPtr = DAG.getCopyFromReg(Chain, DL, Fe::SP,
                                      getPointerTy(DAG.getDataLayout()));
      }
      unsigned LocMemOffset = VA.getLocMemOffset(); //偏移量
      SDValue PtrOff = DAG.getIntPtrConstant(LocMemOffset, DL);
      PtrOff = DAG.getNode(ISD::ADD, DL, getPointerTy(DAG.getDataLayout()),
                           StackPtr, PtrOff);//StackPtr是基址
      /// store val -> reg + offset
      Chain = DAG.getStore(Chain, DL, OutVals[i]/*实参*/, PtrOff,
                           MachinePointerInfo::getStack(MF, LocMemOffset));
    }
  }

  /// 没有参数
  /// 3. 生成Call节点
  GlobalAddressSDNode *N = dyn_cast<GlobalAddressSDNode>(Callee);
  MVT Ty = getPointerTy(DAG.getDataLayout());
  Callee = DAG.getTargetGlobalAddress(N->getGlobal(), DL, Ty);

  SmallVector<SDValue, 8> Ops(1, Chain);
  Ops.push_back(Callee);

  APInt ConstValue;
  std::vector<int> numbers;
  SDValue Glue;
  for (int i = 0, e = RegsPairs.size(); i != e; ++i) {
    auto &[reg, val] = RegsPairs[i];
    Chain = DAG.getCopyToReg(Chain, DL, reg, val, Glue);
    Glue = Chain.getValue(1);
    Ops.push_back(DAG.getRegister(reg, val.getValueType()));
  }

  const TargetRegisterInfo *TRI = Subtarget.getRegisterInfo();
  const uint32_t *Mask = TRI->getCallPreservedMask(DAG.getMachineFunction(), CallConv);
  Ops.push_back(DAG.getRegisterMask(Mask));

  if (Glue.getNode()) {
    Ops.push_back(Glue);
  }

  SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
  Chain = DAG.getNode(FeISD::Call, DL, NodeTys, Ops);

  {
    /// 4. 处理Call的返回值，根据Ins，填充InVals
    SDValue Glue = Chain.getValue(1);
    SmallVector<CCValAssign, 2> RVLos;
    CCState CCInfo(CallConv, IsVarArg, DAG.getMachineFunction(), RVLos,

                   *DAG.getContext());
    CCInfo.AnalyzeCallResult(Ins, RetCC_Fe);
    for (unsigned i = 0, e = RVLos.size(); i != e; ++i) {
      CCValAssign &VA = RVLos[i];
      EVT vt = RVLos[i].getLocVT();
      assert(VA.isRegLoc());
      unsigned RVReg = VA.getLocReg();
      SDValue Val = DAG.getCopyFromReg(Chain, DL, RVReg, vt, Glue);
      Chain = Val.getValue(1);
      Glue = Val.getValue(2);
      InVals.push_back(Val);
    }
  }
//  std::string str(N->getGlobal()->getName());
//  dataCtrl.list.append(str,numbers);
//  if(str == "DIN"){
//    printf("DIN\n");
//    int x0 = numbers.at(0);
//    int x1 = 1024/32;                  //32位
//    dataCtrl.displaymemData();
//    dataCtrl.insdin(x0,x1);
//    dataCtrl.displayInbufferData();
//  }
//  if(str=="DOUT"){
//    printf("DOUT\n");
//    int x0 = numbers.at(0);
//    int x1 = numbers.at(1);
//    dataCtrl.displayOutbufferData();
//    dataCtrl.insdout(x0,x1);
//    dataCtrl.displaymemData();
//  }
//  if(str =="RAC"){
//    printf("RAC\n");
//    int x0 = numbers.at(0);
//    int x1 = numbers.at(1);
//    dataCtrl.insrac(x0,x1);
//    printf("RAC = %d\n",dataCtrl.RAC);
//  }
  return Chain;
}

const char *FeTargetLowering::getTargetNodeName(unsigned Opcode) const {
  switch (Opcode) {
  case FeISD::RET_GLUE:
    return "FeISD::RET_GLUE";
  case FeISD::STO:
    return "FeISD::STO";
  case FeISD::LOA:
    return "FeISD::LOA";
  case FeISD::Call:
    return "FeISD::Call";
  default:
    return nullptr;
  }
}